import torch
import torchvision
from torchvision.transforms import transforms
from torchvision.datasets import ImageFolder
from torchvision.models import DenseNet
from torch.utils.data import Dataset, DataLoader
import torch.nn as nn
from torchvision.transforms import ToTensor, Resize, Compose
import torch.nn.functional as F


feature_extracting = True


def set_parameter_requires_grad(model, feature_extracting):
    if feature_extracting:
        for param in model.parameters():
            param.requires_grad = False


def initialize_model(model_name, num_classes, feature_extract, use_pretrained=True):

    model_fit = None
    if model_name == "DenseNet":
        model_fit = torchvision.models.DenseNet()
        set_parameter_requires_grad(model_fit, feature_extracting=feature_extract)

        dim_fc = model_fit.classifier.in_features

        model_fit.classifier = nn.Sequential(
         nn.Linear(dim_fc, dim_fc, ),
        nn.Linear(dim_fc, num_classes)
        )

    return model_fit


mymodel = initialize_model(model_name="DenseNet", num_classes=4, feature_extract=True)

img = torch.randn((1,3, 224, 224))

y = torch.tensor([1])

loss_func = nn.CrossEntropyLoss()

optimizer = torch.optim.SGD(mymodel.parameters(), lr=0.001)


predict = mymodel(img)
loss = loss_func(predict, y)
loss.backward()
print(loss.data)
optimizer.step()
